from langchain_core.output_parsers import StrOutputParser
from langchain_core.output_parsers.openai_tools import JsonOutputToolsParser
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
# 加载 .env 到环境变量
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())


# 方法里面必须要描述才能被调用，比如下面的  """两个整数相乘"""
@tool
def multiply(first_int: int, second_int: int) -> int:
    """两个整数相乘"""
    return first_int * second_int


@tool
def add(first_int: int, second_int: int) -> int:
    "Add two integers."
    return first_int + second_int


@tool
def exponentiate(base: int, exponent: int) -> int:
    "Exponentiate the base to the exponent power."
    return base ** exponent


tools = [multiply, add, exponentiate]
# 模型
model = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
# 带有分支的 LCEL
llm_with_tools = model.bind_tools(tools) | {
    "functions": JsonOutputToolsParser(),
    "text": StrOutputParser()
}

result = llm_with_tools.invoke("1024的6倍是多少")
print(result)
